import xlwt
from exit import db


def sql_output_excel(sql, filename, sheet_name):
    title = db.session.execute(sql).keys()
    data = db.session.execute(sql).fetchall()

    output_excel(title, data, filename, sheet_name)


def output_excel(title, data, filename, sheet_name):

    book = xlwt.Workbook(encoding='utf-8')

    # 设置颜色
    pattern = xlwt.Pattern()
    pattern.pattern = xlwt.Pattern.SOLID_PATTERN
    pattern.pattern_fore_colour = 22

    # 设置边框
    borders = xlwt.Borders()
    borders.right = xlwt.Borders.THIN
    borders.top = xlwt.Borders.THIN
    borders.bottom = xlwt.Borders.THIN
    borders.left = xlwt.Borders.THIN

    style = xlwt.XFStyle()
    style.pattern = pattern
    style.borders = borders

    sheet1 = book.add_sheet(sheet_name)

    # 写入头部
    col_num = 0
    for single_data in title:
        sheet1.write(0, col_num, single_data, style)
        col_num += 1

    # 写数据
    row_num = 1
    for row_data in data:
        col_num = 0
        for single_data in row_data:
            sheet1.write(row_num, col_num, str(single_data))
            col_num += 1
        row_num += 1

    book.save(filename)

